#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SQLite到MySQL数据库迁移脚本
功能：将项目中的SQLite数据库完整迁移到MySQL数据库
"""

import os
import sys
import sqlite3
import pymysql
import pandas as pd
from sqlalchemy import create_engine
from sqlalchemy.exc import SQLAlchemyError

# MySQL连接配置
MYSQL_CONFIG = {
    'user': 'root',
    'password': '123456',
    'host': 'localhost',
    'port': 3306,
    'database': 'company'
}

# SQLite数据库路径
SQLITE_DB_PATH = os.path.join(os.path.dirname(__file__), 'enterprise.db')

# 表名映射 - 确保与项目中的模型一致
TABLES = [
    'users',
    'departments',
    'projects',
    'tasks',
    'attendance_records'
]


def create_mysql_database():
    """创建MySQL数据库"""
    try:
        # 连接到MySQL服务器（不指定数据库）
        conn = pymysql.connect(
            user=MYSQL_CONFIG['user'],
            password=MYSQL_CONFIG['password'],
            host=MYSQL_CONFIG['host'],
            port=MYSQL_CONFIG['port'],
            charset='utf8mb4'
        )
        cursor = conn.cursor()
        
        # 创建数据库（如果不存在）
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {MYSQL_CONFIG['database']} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;")
        print(f"✅ 已创建MySQL数据库：{MYSQL_CONFIG['database']}")
        
        cursor.close()
        conn.close()
        return True
    except Exception as e:
        print(f"❌ 创建MySQL数据库失败：{e}")
        return False


def get_sqlite_tables():
    """获取SQLite中所有表名"""
    try:
        conn = sqlite3.connect(SQLITE_DB_PATH)
        cursor = conn.cursor()
        
        # 查询所有表名
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';")
        tables = [table[0] for table in cursor.fetchall()]
        
        cursor.close()
        conn.close()
        return tables
    except Exception as e:
        print(f"❌ 获取SQLite表结构失败：{e}")
        return []


def migrate_table_data(table_name):
    """迁移单个表的数据"""
    try:
        # 连接到SQLite数据库
        sqlite_conn = sqlite3.connect(SQLITE_DB_PATH)
        
        # 读取SQLite表数据
        df = pd.read_sql_query(f"SELECT * FROM {table_name};", sqlite_conn)
        print(f"📊 读取表 {table_name}，共 {len(df)} 条记录")
        
        # 关闭SQLite连接
        sqlite_conn.close()
        
        if len(df) == 0:
            print(f"⏭️  表 {table_name} 没有数据，跳过")
            return True
        
        # 构建MySQL连接字符串
        mysql_url = f"mysql+pymysql://{MYSQL_CONFIG['user']}:{MYSQL_CONFIG['password']}@{MYSQL_CONFIG['host']}:{MYSQL_CONFIG['port']}/{MYSQL_CONFIG['database']}?charset=utf8mb4"
        
        # 连接到MySQL并插入数据
        mysql_engine = create_engine(mysql_url)
        
        # 写入MySQL表，替换现有表（如果存在）
        df.to_sql(
            name=table_name,
            con=mysql_engine,
            if_exists='replace',
            index=False,
            chunksize=1000
        )
        
        print(f"✅ 成功迁移表 {table_name} 数据到MySQL")
        return True
    except Exception as e:
        print(f"❌ 迁移表 {table_name} 数据失败：{e}")
        return False


def update_project_config():
    """更新项目配置文件，将数据库改为MySQL"""
    try:
        # 更新主目录下的.env文件
        env_path = os.path.join(os.path.dirname(__file__), '.env')
        new_content = f"DATABASE_URL=mysql+pymysql://{MYSQL_CONFIG['user']}:{MYSQL_CONFIG['password']}@{MYSQL_CONFIG['host']}:{MYSQL_CONFIG['port']}/{MYSQL_CONFIG['database']}\n"
        
        # 读取现有.env内容并更新
        if os.path.exists(env_path):
            with open(env_path, 'r', encoding='utf-8') as f:
                content_lines = f.readlines()
                
            # 查找并替换DATABASE_URL行
            updated = False
            for i, line in enumerate(content_lines):
                if line.startswith('DATABASE_URL='):
                    content_lines[i] = new_content
                    updated = True
                    break
            
            # 如果没有找到，就添加到文件末尾
            if not updated:
                content_lines.append('\n' + new_content)
            
            # 写回文件
            with open(env_path, 'w', encoding='utf-8') as f:
                f.writelines(content_lines)
        else:
            # 如果文件不存在，创建新文件
            with open(env_path, 'w', encoding='utf-8') as f:
                f.write(new_content)
        
        print(f"✅ 已更新项目配置文件：{env_path}")
        
        # 同样更新backend目录下的.env文件
        backend_env_path = os.path.join(os.path.dirname(__file__), 'backend', '.env')
        if os.path.exists(backend_env_path):
            with open(backend_env_path, 'r', encoding='utf-8') as f:
                content_lines = f.readlines()
                
            # 查找并替换DATABASE_URL行
            updated = False
            for i, line in enumerate(content_lines):
                if line.startswith('DATABASE_URL='):
                    content_lines[i] = new_content
                    updated = True
                    break
            
            # 如果没有找到，就添加到文件末尾
            if not updated:
                content_lines.append('\n' + new_content)
            
            # 写回文件
            with open(backend_env_path, 'w', encoding='utf-8') as f:
                f.writelines(content_lines)
            
            print(f"✅ 已更新后端配置文件：{backend_env_path}")
        
        return True
    except Exception as e:
        print(f"❌ 更新项目配置文件失败：{e}")
        return False


def setup_admin_system():
    """确保管理员系统正确设置"""
    try:
        # 调用项目中的setup_admin_system.py脚本
        setup_script = os.path.join(os.path.dirname(__file__), 'setup_admin_system.py')
        if os.path.exists(setup_script):
            print("🔧 正在执行管理员系统设置...")
            os.system(f'python "{setup_script}"')
            print("✅ 管理员系统设置完成")
        else:
            print("⚠️  setup_admin_system.py 脚本不存在，跳过管理员设置")
        return True
    except Exception as e:
        print(f"❌ 设置管理员系统失败：{e}")
        return False


def main():
    """主函数"""
    print("🚀 开始SQLite到MySQL数据库迁移...")
    print(f"📌 SQLite数据库路径: {SQLITE_DB_PATH}")
    print(f"📌 MySQL数据库配置: {MYSQL_CONFIG['user']}@{MYSQL_CONFIG['host']}:{MYSQL_CONFIG['port']}/{MYSQL_CONFIG['database']}")
    
    # 1. 检查SQLite数据库是否存在
    if not os.path.exists(SQLITE_DB_PATH):
        print("❌ SQLite数据库文件不存在，请检查路径")
        sys.exit(1)
    
    # 2. 安装必要依赖
    try:
        import pymysql
        import pandas
        import sqlalchemy
        print("✅ 依赖库已安装")
    except ImportError:
        print("⚠️  正在安装必要的依赖库...")
        os.system('pip install pymysql pandas sqlalchemy')
    
    # 3. 创建MySQL数据库
    if not create_mysql_database():
        sys.exit(1)
    
    # 4. 获取SQLite中的表
    sqlite_tables = get_sqlite_tables()
    if not sqlite_tables:
        print("⚠️  在SQLite数据库中未找到任何表，可能是空数据库")
        # 继续执行，以便创建空表结构
    else:
        print(f"📋 在SQLite中找到 {len(sqlite_tables)} 个表：{', '.join(sqlite_tables)}")
    
    # 5. 由于SQLite数据库为空，跳过数据迁移步骤
    print("ℹ️  SQLite数据库为空，跳过数据迁移步骤")
    
    # 但是我们仍然需要确保MySQL数据库创建了必要的表结构
    # 这将通过setup_admin_system或后续的应用运行来完成
    
    # 6. 更新项目配置文件
    if not update_project_config():
        print("❌ 配置更新失败，已停止")
        sys.exit(1)
    
    # 7. 安装PyMySQL到项目依赖
    print("🔧 安装PyMySQL到项目依赖...")
    os.system('pip install PyMySQL')
    
    # 8. 运行setup_admin_system确保管理员用户正确设置
    setup_admin_system()
    
    print("\n🎉 数据库迁移完成！")
    print("📝 下一步操作：")
    print("   1. 启动后端服务：cd backend && python -m uvicorn app.main:app --reload")
    print("   2. 启动前端服务：cd frontend && python -m http.server 8080")
    print("   3. 访问 http://localhost:8080 验证迁移是否成功")


if __name__ == "__main__":
    main()